#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
评估数据清理工具
用于删除量表评估数据
"""

import pymysql
import sys
from datetime import datetime

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'port': 3306,
    'user': 'flask_dev',
    'password': 'dev123456',
    'database': 'flask_local',
    'charset': 'utf8mb4'
}

def get_connection():
    """获取数据库连接"""
    try:
        connection = pymysql.connect(**DB_CONFIG)
        return connection
    except Exception as e:
        print(f"❌ 数据库连接失败: {e}")
        sys.exit(1)

def show_statistics(connection):
    """显示当前数据统计"""
    try:
        with connection.cursor() as cursor:
            print("\n" + "="*50)
            print("📊 当前数据统计")
            print("="*50)
            
            # 评估记录统计
            cursor.execute("SELECT COUNT(*) FROM scale_assessment")
            assessment_count = cursor.fetchone()[0]
            print(f"评估记录数: {assessment_count}")
            
            # 量表详细记录统计
            cursor.execute("SELECT COUNT(*) FROM scale_hamd17")
            hamd_count = cursor.fetchone()[0]
            print(f"HAMD详细记录数: {hamd_count}")
            
            cursor.execute("SELECT COUNT(*) FROM scale_hama")
            hama_count = cursor.fetchone()[0]
            print(f"HAMA详细记录数: {hama_count}")
            
            cursor.execute("SELECT COUNT(*) FROM scale_ymrs")
            ymrs_count = cursor.fetchone()[0]
            print(f"YMRS详细记录数: {ymrs_count}")
            
            # 患者统计
            cursor.execute("""
                SELECT COUNT(DISTINCT user_id) 
                FROM scale_assessment
            """)
            patient_count = cursor.fetchone()[0]
            print(f"涉及患者数: {patient_count}")
            
            print("="*50 + "\n")
            
            return assessment_count > 0
    except Exception as e:
        print(f"❌ 查询统计失败: {e}")
        return False

def clear_all_data(connection):
    """清空所有评估数据"""
    try:
        with connection.cursor() as cursor:
            print("⚠️  正在清空所有评估数据...")
            
            # 禁用外键检查
            cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
            
            # 删除所有数据
            cursor.execute("DELETE FROM scale_hamd17")
            hamd_deleted = cursor.rowcount
            
            cursor.execute("DELETE FROM scale_hama")
            hama_deleted = cursor.rowcount
            
            cursor.execute("DELETE FROM scale_ymrs")
            ymrs_deleted = cursor.rowcount
            
            cursor.execute("DELETE FROM scale_assessment")
            assessment_deleted = cursor.rowcount
            
            # 重置自增ID
            cursor.execute("ALTER TABLE scale_assessment AUTO_INCREMENT = 1")
            cursor.execute("ALTER TABLE scale_hamd17 AUTO_INCREMENT = 1")
            cursor.execute("ALTER TABLE scale_hama AUTO_INCREMENT = 1")
            cursor.execute("ALTER TABLE scale_ymrs AUTO_INCREMENT = 1")
            
            # 启用外键检查
            cursor.execute("SET FOREIGN_KEY_CHECKS = 1")
            
            # 提交事务
            connection.commit()
            
            print("\n✅ 清理完成！")
            print(f"   - 删除评估记录: {assessment_deleted} 条")
            print(f"   - 删除HAMD记录: {hamd_deleted} 条")
            print(f"   - 删除HAMA记录: {hama_deleted} 条")
            print(f"   - 删除YMRS记录: {ymrs_deleted} 条")
            print(f"   - 已重置所有表的自增ID")
            
            return True
    except Exception as e:
        connection.rollback()
        print(f"❌ 清理失败: {e}")
        return False

def clear_patient_data(connection, patient_id):
    """清空指定患者的评估数据"""
    try:
        with connection.cursor() as cursor:
            # 查询患者信息
            cursor.execute("""
                SELECT real_name, phone, role 
                FROM user 
                WHERE id = %s
            """, (patient_id,))
            
            patient = cursor.fetchone()
            if not patient:
                print(f"❌ 患者ID={patient_id} 不存在")
                return False
            
            print(f"\n患者信息: 姓名={patient[0]}, 手机={patient[1]}, 角色={patient[2]}")
            
            # 查询该患者的评估记录数
            cursor.execute("""
                SELECT COUNT(*) 
                FROM scale_assessment 
                WHERE user_id = %s
            """, (patient_id,))
            
            count = cursor.fetchone()[0]
            if count == 0:
                print(f"⚠️  患者ID={patient_id} 没有评估记录")
                return True
            
            print(f"该患者有 {count} 条评估记录")
            
            # 删除评估记录（会自动级联删除详细数据）
            cursor.execute("""
                DELETE FROM scale_assessment 
                WHERE user_id = %s
            """, (patient_id,))
            
            deleted = cursor.rowcount
            connection.commit()
            
            print(f"\n✅ 清理完成！删除了 {deleted} 条评估记录")
            return True
    except Exception as e:
        connection.rollback()
        print(f"❌ 清理失败: {e}")
        return False

def backup_database():
    """备份数据库"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = f"backup_before_clear_{timestamp}.sql"
    
    print(f"\n📦 正在备份数据库到: {backup_file}")
    
    import subprocess
    try:
        cmd = [
            'mysqldump',
            '-u', DB_CONFIG['user'],
            f"-p{DB_CONFIG['password']}",
            '-h', DB_CONFIG['host'],
            '--port', str(DB_CONFIG['port']),
            DB_CONFIG['database'],
            'scale_assessment',
            'scale_hamd17',
            'scale_hama',
            'scale_ymrs'
        ]
        
        with open(backup_file, 'w') as f:
            subprocess.run(cmd, stdout=f, check=True, stderr=subprocess.PIPE)
        
        print(f"✅ 备份完成: {backup_file}")
        return True
    except FileNotFoundError:
        print("⚠️  未找到 mysqldump 命令，跳过备份")
        print("   建议手动备份数据库")
        return False
    except Exception as e:
        print(f"❌ 备份失败: {e}")
        return False

def main():
    """主函数"""
    print("\n" + "="*50)
    print("🗑️  评估数据清理工具")
    print("="*50)
    
    # 获取数据库连接
    connection = get_connection()
    
    try:
        # 显示当前统计
        has_data = show_statistics(connection)
        
        if not has_data:
            print("ℹ️  数据库中没有评估数据")
            return
        
        # 选择清理方式
        print("请选择清理方式：")
        print("1. 清空所有评估数据")
        print("2. 清空指定患者的评估数据")
        print("0. 退出")
        
        choice = input("\n请输入选项 (0-2): ").strip()
        
        if choice == '0':
            print("已取消操作")
            return
        
        elif choice == '1':
            # 确认操作
            print("\n⚠️  警告：此操作将删除所有评估数据，不可恢复！")
            confirm = input("是否继续？(yes/no): ").strip().lower()
            
            if confirm != 'yes':
                print("已取消操作")
                return
            
            # 询问是否备份
            backup = input("是否先备份数据？(yes/no): ").strip().lower()
            if backup == 'yes':
                backup_database()
            
            # 清空所有数据
            if clear_all_data(connection):
                # 显示清理后的统计
                show_statistics(connection)
        
        elif choice == '2':
            # 输入患者ID
            patient_id = input("\n请输入患者ID: ").strip()
            
            try:
                patient_id = int(patient_id)
            except ValueError:
                print("❌ 患者ID必须是数字")
                return
            
            # 确认操作
            confirm = input(f"确认删除患者ID={patient_id}的所有评估数据？(yes/no): ").strip().lower()
            
            if confirm != 'yes':
                print("已取消操作")
                return
            
            # 清空指定患者数据
            clear_patient_data(connection, patient_id)
        
        else:
            print("❌ 无效的选项")
    
    finally:
        connection.close()
        print("\n👋 感谢使用！\n")

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠️  操作已取消")
        sys.exit(0)

